from flask import Flask, render_template, request, jsonify
from PIL import Image
import io
import base64

app = Flask(__name__)

# 用于存储用户上传的图片和秘密信息文件
carrier_image = None
secret_image = None

def plus(value):
    return bin(value)[2:].zfill(8)


def getCode(img):
    code = ""
    for i in range(img.size[0]):
        for j in range(img.size[1]):
            gray_value = img.getpixel((i, j))
            code += plus(gray_value)  # 使用 gray_value，无需将其转换为整数

    return code

def encry(img, code):
    count = 0
    codeLen = len(code)

    for i in range(img.size[0]):
        for j in range(img.size[1]):
            gray_value = img.getpixel((i, j))  # 提取元组中的灰度值

            if count == codeLen:
                break

            new_gray_value = (gray_value - gray_value % 2) + int(code[count])

            count += 1
            if count == codeLen:
                img.putpixel((i, j), new_gray_value)
                break

            img.putpixel((i, j), new_gray_value)
    return img

def deEncry(img, length):
    width = img.size[0]
    height = img.size[1]
    count = 0
    wt = ""

    for i in range(width):
        for j in range(height):
            gray_value = img.getpixel((i, j))  # 提取元组中的整数灰度值
            wt = wt + str(gray_value % 2)
            count += 1

            if count == length:
                break

        if count == length:
            break

    return wt
def showImage(wt, img_width, img_height):
    str1 = []
    for i in range(0, len(wt), 8):
        # 以每8位为一组二进制，转换为十进制
        str1.append(int(wt[i:i + 8], 2))

    # 图片大于水印图片1个像素，便于对比
    img_out = Image.new("L", (img_width, img_height))  # 使用灰度模式
    flag = 0
    for m in range(0, img_width):
        for n in range(0, img_height):
            if flag == len(str1):
                break
            img_out.putpixel((m, n), str1[flag])
            flag += 1
        if flag == len(str1):
            break
    return img_out



@app.route('/')
def index():
    return render_template('Display_page.html')

@app.route('/upload_carrier', methods=['POST'])
def upload_carrier():
    global carrier_image
    try:
        if 'file' in request.files:
            file = request.files['file']
            if file:
                # 检查文件类型
                if file.filename.endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                    image = Image.open(io.BytesIO(file.read()))
                    carrier_image = image
                    return jsonify({'success': True})
                else:
                    return jsonify({'success': False, 'error': 'Invalid file format'})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

    return jsonify({'success': False})


@app.route('/upload_secret', methods=['POST'])
def upload_secret():
    global secret_image
    try:
        if 'file' in request.files:
            file = request.files['file']
            if file:
                # 检查文件类型
                if file.filename.endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                    image = Image.open(io.BytesIO(file.read()))
                    secret_image = image
                    return jsonify({'success': True})
                else:
                    return jsonify({'success': False, 'error': 'Invalid file format'})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

    return jsonify({'success': False})



@app.route('/embed_watermark', methods=['POST'])
def embed_watermark_endpoint():
    global carrier_image, secret_image, gray_img1 # 添加 gray_img1
    if carrier_image is not None and secret_image is not None:
        gray_img1 = carrier_image.convert('L')
        gray_img2 = secret_image.convert('L')
        secret_code = getCode(gray_img2)
        encry(gray_img1, secret_code)
        buffered = io.BytesIO()
        gray_img1.save(buffered, format="bmp")
        result_image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
        return jsonify({'success': True, 'result': result_image_base64})
    return jsonify({'success': False})


@app.route('/extract_watermark', methods=['POST'])
def extract_watermark_endpoint():
    global secret_image,gray_img1
    if secret_image is not None:
        image_width = secret_image.size[0]
        image_height = secret_image.size[1]
        length = image_width * image_height * 8
        gray_img1 = gray_img1.convert('L')  # 使用灰度模式
        wt = deEncry(gray_img1, length)
        showImage(wt,image_width,image_height)
        image = showImage(wt,image_width,image_height)
        buffered = io.BytesIO()
        image.save(buffered, format="bmp")
        result_image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
        return jsonify({'success': True, 'result': result_image_base64})
    return jsonify({'success': False})


if __name__ == '__main__':
    app.run(debug=False, use_reloader=False, use_debugger=False)



